package inference.traffic;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.NavigableMap;
import java.util.NavigableSet;
import java.util.TreeMap;
import java.util.TreeSet;

import com.sun.xml.internal.ws.api.pipe.NextAction;

import util.Pair;

public class ProblemGenerator {

	private Pair<Integer, Integer> bottomLeft, bottomRight, topLeft, topRight;
	
	private NavigableSet< Pair<Integer, Integer > > _occupied_cells = new TreeSet<Pair<Integer,Integer>>();
	
	private NavigableSet< Pair<Integer, Integer> > _input_cells = new TreeSet<Pair<Integer,Integer>>();
	
	private NavigableMap< Pair<Integer,Integer>,Pair<Boolean,Boolean> > _light_signals 
		= new TreeMap<Pair<Integer,Integer>, Pair<Boolean,Boolean>>();
	
	private NavigableSet< Pair< Integer, Integer> > _exit_cells = new TreeSet<Pair<Integer,Integer>>();

	private NavigableMap< Pair< Integer, Integer >, Boolean > _turn_dirxn 
		= new TreeMap<Pair<Integer,Integer>, Boolean>();
	
	private NavigableMap< Pair<Integer, Integer>, Pair<Integer, Integer> > _nextCell
		= new TreeMap< Pair<Integer, Integer>, Pair<Integer, Integer> >();
	
	private NavigableMap< Pair<Integer, Integer>, Pair<Integer, Integer> > _prevCell
	= new TreeMap< Pair<Integer, Integer>, Pair<Integer, Integer> >();
	
	private NavigableMap< Pair<Integer, Integer>, Pair<Boolean,Boolean> > _xion_dirxn
		= new TreeMap<Pair<Integer,Integer>, Pair<Boolean,Boolean>>();

	private Random _rand;
	
	private final String _instance_name;
	
	public NavigableMap<Pair<Integer, Integer>, Pair<Boolean, Boolean>> get_xion_dirxn() {
		return _xion_dirxn;
	}
	
	public NavigableSet<Pair<Integer, Integer>> get_occupied_cells() {
		return _occupied_cells;
	}

	public NavigableSet<Pair<Integer, Integer>> get_input_cells() {
		return _input_cells;
	}

	public NavigableMap<Pair<Integer, Integer>, Pair<Boolean, Boolean>> get_light_signals() {
		return _light_signals;
	}

	public NavigableSet<Pair<Integer, Integer>> get_exit_cells() {
		return _exit_cells;
	}

	public NavigableMap<Pair<Integer, Integer>, Boolean> get_turn_dirxn() {
		return _turn_dirxn;
	}

	public NavigableMap<Pair<Integer, Integer>, Pair<Integer, Integer>> get_nextCell() {
		return _nextCell;
	}

	public NavigableMap<Pair<Integer, Integer>, Pair<Integer, Integer>> get_prevCell() {
		return _prevCell;
	}
	
	public String get_instance_name() {
		return _instance_name;
	}

	public ProblemGenerator( final int nrows, final int ncols, final int wingSize,
								final int blockSize, final double turnProb,
								final double arrivalProb, int horizon, double discount,
								final long rand_seed ){
		
		_rand = new Random( rand_seed );
		
		_instance_name = "traffic_mdp_" + (nrows) + "_" + (ncols);
		
		int maxx_grid = wingSize + (ncols-1) + (ncols-1)*blockSize;
		int maxy_grid = wingSize + (nrows-1) + (nrows-1)*blockSize ;
		
		bottomLeft = new Pair<Integer, Integer>(wingSize, wingSize);
		bottomRight = new Pair<Integer, Integer>( maxx_grid, wingSize);

		topLeft = new Pair<Integer, Integer>( wingSize, maxy_grid );
		topRight = new Pair<Integer, Integer>( maxx_grid, maxy_grid );

		int MAXX = maxx_grid + wingSize;
		int MAXY = maxy_grid + wingSize;
		
		StringBuffer non_fluents = new StringBuffer();
		
		NavigableSet<String> cells = new TreeSet<String>();

		NavigableSet<String> intersections = new TreeSet<String>();
		
		NavigableMap< String, String> cell_xion = new TreeMap<String, String>();
		
		NavigableSet<String> street_dirxn = new TreeSet<String>();
		
		NavigableSet<String> input_cells = new TreeSet<String>();
		NavigableSet<String> exit_cells = new TreeSet<String>();
		
		NavigableMap<String, String> flows_x = new TreeMap<String, String>();
		NavigableMap<String, String> flows_y = new TreeMap<String, String>();
		
		NavigableMap<String, Set<String> > adjacency = new TreeMap<String, Set<String>>();

		int mod_factor = (blockSize+1);
		
		for( int i = 0 ; i <= MAXX; ++i ){
			for( int j = 0 ; j <= MAXY; ++j ){
				//is this a cell?
				
				if( (i >= bottomLeft._o1 && i <= bottomRight._o1 && (i-wingSize)%mod_factor == 0 ) || 
						( j >= bottomLeft._o2 && j <= topLeft._o2 && (j-wingSize)%mod_factor == 0 ) ){
					
					String cellName = "c"+(i)+"_"+(j);
					
					cells.add( cellName );
					
					//this an x street?
					if(  (j >= bottomLeft._o2 && j <= topRight._o2 )
							&& ( j - wingSize )%mod_factor == 0 ){
					
						if( ((int)((j-wingSize)/mod_factor)) %2 == 0 ){
							street_dirxn.add("STREET-EW(" + cellName + ")");
						}else{
							street_dirxn.add("STREET-WE(" + cellName + ")");
						}
						
					}

					//is this a y street
					if(  ( i >= bottomLeft._o2 && i <= topLeft._o2 ) &&
							 ( i - wingSize )%mod_factor == 0 ){
					
						if( ((int)((i-wingSize)/mod_factor))%2 == 0 ){
							street_dirxn.add("STREET-NS(" + cellName + ")");
						}else{
							street_dirxn.add("STREET-SN("+ cellName + ")");
						}
						
					}
					
					//is this an xion?
					if( (i >= bottomLeft._o1 && i <= bottomRight._o1) && 
							( j >= bottomLeft._o2 && j <= topLeft._o2 ) &&
							( i - wingSize )% mod_factor == 0 &&
							( j - wingSize )% mod_factor == 0 ){
						
						String xion_name = "i" + (i) + "_" + (j);
						
						intersections.add( xion_name );
						
						_light_signals.put( new Pair<Integer, Integer>(i,j), 
								new Pair<Boolean, Boolean>(false, false) );

						_turn_dirxn.put( new Pair<Integer, Integer>(i,j),
								TrafficState.getBernoulli( _rand, turnProb ) );
						
						boolean one, two;
						
						if( street_dirxn.contains("STREET-EW(" + cellName + ")") ){
							one = true;
						}else{
							one = false;
						}

						if( street_dirxn.contains("STREET-NS(" + cellName + ")" ) ){
							two = true;
						}else{
							two = false;
						}
						
						_xion_dirxn.put(new Pair<Integer, Integer>(i,j), new Pair<Boolean, Boolean>(one, two) );
						
						
						cell_xion.put( cellName, xion_name );
						
						if( i-mod_factor >= 0 ){
							
							if( adjacency.get(xion_name ) == null ){
								adjacency.put( xion_name , new TreeSet<String>() );
							}
							
							adjacency.get( xion_name ).add( "i"+(i-mod_factor)+"_"+(j) );
							
						}
						
						if( i+mod_factor <= MAXX ){
							
							if( adjacency.get( xion_name ) == null ){
								adjacency.put( xion_name , new TreeSet<String>() );
							}
							
							adjacency.get( xion_name ).add("i"+(i+mod_factor)+"_" + (j) );
						}
						
						if( j-mod_factor >= 0 ){
							
							if( adjacency.get( xion_name ) == null ){
								adjacency.put( xion_name , new TreeSet<String>() );
							}
							
							adjacency.get( xion_name ).add("i"+(i)+"_" + (j-mod_factor) );
						}
						
						if( j+mod_factor <= MAXY ){
							
							if( adjacency.get( xion_name ) == null ){
								adjacency.put( xion_name , new TreeSet<String>() );
							}
							
							adjacency.get( xion_name ).add("i"+(i)+"_" + (j+mod_factor) );
						}
						
					}

							
					//covers rate as well
					//input cell?
					if( j >= bottomLeft._o2 && j <= topLeft._o2 && 
							(
									( i >= 0 && i < bottomLeft._o1 && street_dirxn.contains( "STREET-WE(" + cellName + ")" ) )
									||
									( i <= MAXX && i > bottomRight._o1 && street_dirxn.contains( "STREET-EW(" + cellName + ")" ) )
							) ){
						
						input_cells.add("PERIMETER-INPUT-CELL("+cellName+")");
						
						_input_cells.add( new Pair<Integer, Integer>( i, j ) );
						
					}
					
					if( i >= bottomLeft._o1 && i <= bottomRight._o1 && 
							(
									( j >= 0 && j < bottomLeft._o2 && street_dirxn.contains( "STREET-SN(" + cellName + ")" ) )
									||
									( j <= MAXY && j > topLeft._o2 && street_dirxn.contains( "STREET-NS(" + cellName + ")" ) )
							) ){
						
						input_cells.add("PERIMETER-INPUT-CELL("+ cellName +")");
						
						_input_cells.add( new Pair<Integer, Integer>( i, j ) );
						
					}
					
					//exit cell?
					if( j >= bottomLeft._o2 && j <= topLeft._o2 && 
							(
									( i == 0 && street_dirxn.contains( "STREET-EW(" + cellName + ")" ) )
									||
									( i == MAXX && street_dirxn.contains( "STREET-WE(" + cellName + ")" ) )
							) ){
						
						exit_cells.add("PERIMETER-EXIT-CELL("+ cellName +")");
						
						_exit_cells.add( new Pair<Integer, Integer>(i,j) );
						
					}
					
					if( i >= bottomLeft._o1 && i <= bottomRight._o1 && 
							(
									( j == 0 && street_dirxn.contains( "STREET-NS(" + cellName + ")" ) )
									||
									( j == MAXY && street_dirxn.contains( "STREET-SN(" + cellName + ")" ) )
							) ){
						
						exit_cells.add("PERIMETER-EXIT-CELL(" + cellName +")");
						
						_exit_cells.add( new Pair<Integer, Integer>(i,j) );
						
					}
					
					//flows x?
					if( street_dirxn.contains( "STREET-EW(" + cellName + ")" ) &&
							i-1 >= 0 ){
						flows_x.put( cellName, "c"+(i-1)+ "_"+j );
						
						_nextCell.put( new Pair<Integer, Integer>(i, j), new Pair<Integer, Integer>(i-1,j) );
						_prevCell.put(new Pair<Integer, Integer>(i-1, j), new Pair<Integer, Integer>(i, j) );
						
					}else if( street_dirxn.contains( "STREET-WE(" + cellName + ")" ) &&
							i+1 <= MAXX ){
						flows_x.put( cellName, "c"+(i+1)+"_"+(j) );
						
						_nextCell.put( new Pair<Integer, Integer>(i, j), new Pair<Integer, Integer>(i+1,j) );
						_prevCell.put(new Pair<Integer, Integer>(i+1, j), new Pair<Integer, Integer>(i, j) );
						
					}
					
					//flows y?
					if( street_dirxn.contains( "STREET-SN(" + cellName + ")" ) &&
							j+1 <= MAXX ){
						flows_y.put( cellName, "c"+(i)+"_"+(j+1) );
						
						_nextCell.put( new Pair<Integer, Integer>(i, j), new Pair<Integer, Integer>(i,j+1) );
						_prevCell.put(new Pair<Integer, Integer>(i, j+1), new Pair<Integer, Integer>(i, j) );
						
						
					}else if( street_dirxn.contains( "STREET-NS(" + cellName + ")" ) &&
							j-1 >= 0 ){
						flows_y.put( cellName, "c"+(i)+"_"+(j-1) );

						_nextCell.put( new Pair<Integer, Integer>(i, j), new Pair<Integer, Integer>(i,j-1) );
						_prevCell.put(new Pair<Integer, Integer>(i, j-1), new Pair<Integer, Integer>(i, j) );
						
						
					}
					
					
				}
			}
		}
		
		
		non_fluents.append("non-fluents nf_"+_instance_name + " {\n");
		non_fluents.append("domain = traffic_mdp;\n");
		non_fluents.append("objects {\n");

		non_fluents.append("cell : { ");
		
		int i = 0;
		for( String c : cells ){
			non_fluents.append( c + (( i == cells.size()-1 ) ?  "};\n" : ", " ) );
			++i;
		}
		
		i = 0;

		non_fluents.append("intersection : { ");
		for( String x :  intersections ){
			non_fluents.append( x + (( i == intersections.size()-1 ) ?  "};\n" : ", " ) );
			++i;
		}
		
		non_fluents.append("};\n\n");

		non_fluents.append("non-fluents {\n");

		i = 0;
		
		for( String per : input_cells ){
			non_fluents.append( per +  " = true;\n"  );
			++i;
		}
		
		non_fluents.append("\n");
		
		for( String per : input_cells ){
			String copy = new String(per);
			copy = copy.replaceAll("CELL", "RATE");
			
			non_fluents.append( copy + " = " + arrivalProb + ";\n" );
		}
		
		non_fluents.append("\n");
		
		for( String ex : exit_cells ){
			non_fluents.append( ex + " = true;\n");
		}
		
		for( String x : intersections ){
			non_fluents.append("TURN-PROB(" + x + ") = " + turnProb + ";\n" );
		}
		
		non_fluents.append("\n");
		
		for( Map.Entry<String, String> entry : flows_x.entrySet() ){
			
			non_fluents.append("FLOWS-INTO-X(" + entry.getKey() + " , " + entry.getValue() + ") = true;\n" );
			
		}
		
		for( Map.Entry<String, String> entry : flows_y.entrySet() ){
			
			non_fluents.append("FLOWS-INTO-Y(" + entry.getKey() + " , " + entry.getValue() + ") = true;\n" );
			
		}
		
		for( Map.Entry<String, String> entry : cell_xion.entrySet() ){
			non_fluents.append("INTERSECTION-AT(" + entry.getValue() + ", " + entry.getKey() + ") = true;\n" );
		}
		non_fluents.append("\n");

		for( String str : street_dirxn ){
			non_fluents.append(str+" = true;\n");
		}
		non_fluents.append("\n");
		
		for( Map.Entry<String, Set<String> > entry : adjacency.entrySet() ){

			for( String str : entry.getValue() ){
				
				non_fluents.append("ADJACENT(" + entry.getKey() + ", " + str + ") = true;\n" );
				non_fluents.append("ADJACENT(" + str + ", " + entry.getKey() + ") = true;\n" );
				
			}

		}
		
		non_fluents.append("};\n");
		
		non_fluents.append("\n");
		
		non_fluents.append("}\n\n");
		
		StringBuffer instance = new StringBuffer();
		
		instance.append("instance " + _instance_name + "{\n" );
		
		instance.append("domain = traffic_mdp;\n" );
		
		instance.append("non-fluents = nf_" + _instance_name + ";\n" );
		
		instance.append("\ninit-state {\n");
		
		//legal init state
		for( String cell : cells ){
			
			if( street_dirxn.contains("STREET-EW(" + cell + ")" ) ||
					street_dirxn.contains("STREET-WE(" + cell + ")" ) ){
				instance.append("direction(" + cell + ") = true;\n" );
			}else if( street_dirxn.contains("STREET-SN(" + cell + ")" ) ||
					street_dirxn.contains("STREET-NS(" + cell + ")" ) ){
				instance.append("direction(" + cell + ") = false;\n" );
			}
			
		}
		
		instance.append("\n};\n");
		
		instance.append("max-nondef-actions = " + 2*intersections.size() + ";\n");
		
		instance.append("horizon = " + horizon + ";\n" );
		
		instance.append("discount = " + discount + ";\n" );
		
		instance.append("}\n");
		
		PrintStream ps;
		
		try {
			
			ps = new PrintStream(
					new FileOutputStream( _instance_name + ".rddl" ) );
			
			ps.println( non_fluents.toString() );
			ps.println( instance.toString() );
			
			ps.close();
			
			System.out.println( "Problem instance written to " + _instance_name + ".rddl" );

			
		} catch (FileNotFoundException e) {
		}
		
	}
	
	public static void main(String[] args) {
		new ProblemGenerator(10, 10, 2, 2, 0.3, 0.3, 1, 0.9, 0);
	}
	
}
